import albumentations as albu
from scipy import ndimage
import numpy as np


def aug_data(image,image_size,p =1,):
    image = albu.Compose([
        albu.CoarseDropout(max_holes=4,max_width=4,max_height=4),
        albu.HorizontalFlip(),
        # albu.VerticalFlip(),
        albu.Downscale(scale_max=0.9,scale_min=0.75),
        albu.GaussNoise(p=0.3),
        albu.GaussianBlur(p=0.2),
        albu.MedianBlur(p=0.3),
        albu.RandomBrightnessContrast(p=0.8),
        albu.ElasticTransform(p=0.3),
        albu.Rotate(limit=45),
        albu.RandomResizedCrop(width=image_size[0],height=image_size[1],scale=(0.3,1),p = 1.),
    ],p=p)(image = image)['image']
    return image


def random_rotate_and_crop(image,heatmap=None,crop_limit=20,p = 1,angle_limit = 15):
    if np.random.uniform() < p:
        ## 输入的是image和heatmap。注意旋转后因为避免损失图片信息，尺寸会变大。
        ## 再通过随机裁剪操作还原成原来的尺寸
        ## 旋转
        ori_size = image.shape
        angle = np.random.randint(0-angle_limit,angle_limit)
        image  = ndimage.rotate(image,angle=angle,axes=(1,0))

        ## 旋转后尺寸必定变大，进行裁剪。为了保持在得到box后裁剪进行预测的尺寸和训练的尺寸一致，则
        ## 根据图片尺寸转换后进行裁剪操作。
        random_cut_x_1 = np.random.randint(1,crop_limit)
        random_cut_x_2 = np.random.randint(1,crop_limit)
        random_cut_y_1 = np.random.randint(1,crop_limit)
        random_cut_y_2 = np.random.randint(1,crop_limit)

        image = image[random_cut_x_1:-random_cut_x_2,random_cut_y_1:-random_cut_y_2,:]
        new_size = image.shape
        image = ndimage.zoom(image, (ori_size[0] / new_size[0], ori_size[1] / new_size[1], 1), order=2)
        if heatmap is not None:
            heatmap = ndimage.rotate(heatmap, angle=angle, axes=(2, 1))
            heatmap = heatmap[:, random_cut_x_1:-random_cut_x_2,random_cut_y_1:-random_cut_y_2]
            heatmap = ndimage.zoom(heatmap,(1,ori_size[0]/new_size[0],ori_size[1]/new_size[1]),order=1)
    if heatmap is not None:
        return image,heatmap
    else:
        return image